#ifndef TERN_MEMOIZER_H
#define TERN_MEMOIZER_H

#include <vector>
#include <string>
#include <map>
#include <set>
#include <fstream>

#include <pthread.h>
#include <semaphore.h>
#include <ext/hash_map>
#include <fcntl.h>
#include <errno.h>
#include <unistd.h>
#include <sys/syscall.h>
#include <time.h>
#include <sys/time.h> 

#include "llvm/Instructions.h"
#include "llvm/Function.h"
#include "llvm/Module.h"
#include "llvm/PassManager.h"

#include "klee/ExecutionState.h"
#include "klee/Internal/Module/KInstruction.h"
#include "tern/TernDefs.h"
#include "tern/TernSched.h"
#include "tern/SchedCache.h"

// Include Jingyue's branch slicing module.
#include "../../llvm-instrument/omit-branch/omit-branch.h"

using namespace llvm;
using namespace klee;
using namespace defens;
using namespace std;

//#define TERN_DEBUG_ORDER
#ifdef TERN_DEBUG_ORDER
#        define order_print(args...) fprintf(stderr, args)
#else
#        define order_print(args...)
#endif

//#define TERN_DEBUG_MEMOIZER
#ifdef TERN_DEBUG_MEMOIZER
#        define mmz_print(args...) fprintf(stderr, args);fflush(stderr)
#else
#        define mmz_print(args...)
#endif

// extern functions.
extern pid_t gettid();

namespace klee {
	enum SyncOpType {
		UNKNOWN = -1,
		SIGNAL = 0,
		BROADCAST = 1
	};

	// This struct maintains all stats of logging mem access (determinism in load/store to shared addr).
	struct LogMemAccess {
		long loadStoreInstCnt;
		long addrCnt;
		char *InstToAddrMap[INST_SIZE];
		bool loggingMemAccess;
	};

	class BlockingTidInfo {
	public:
		BlockingTidInfo(int threadId, uint64_t sVarAddr, uint64_t secondSVarAddr) {
			tid = threadId;
			syncVarAddr = sVarAddr;
			secondSyncVarAddr = secondSVarAddr;
			opType = UNKNOWN;
		}
		void setSyncOpType(SyncOpType sot) {
			opType = sot;
		}
		SyncOpType getSyncOpType() {return opType;}
		int getTid() {return tid;}
		uint64_t getSyncVarAddr() {return syncVarAddr;}
		uint64_t getSecondSyncVarAddr() {return secondSyncVarAddr;}
		void setSyncVarAddr(uint64_t sVarAddr, uint64_t secondSVarAddr) {
			syncVarAddr = sVarAddr;
			secondSyncVarAddr = secondSVarAddr;
		}
	private:
		int tid;
		uint64_t syncVarAddr;
		uint64_t secondSyncVarAddr;
		SyncOpType opType;
	};

	class Queue {
	private:
 		vector<int> activeThreadQueue;
		vector<BlockingTidInfo *> blockingThreadQueue;
		vector<int> deadThreadQueue;
		pthread_cond_t *activeQueueCondVar;
		pthread_cond_t *blockingQueueCondVar;
		Int_To_Int_Map tidMap;
	 	int inTaskThreadNum;
		int outTaskThreadNum;
		int ids;
			
	public:	
		Queue(pthread_cond_t *active, pthread_cond_t *blocking);
		void printDeadQueue(); 
		//bool IsDeadQueueEmpty(); 
		bool isTidInDeadQueue(int pthreadId); 
		bool delPthreadIDFromDead(int pthreadId); 
		void addMyselfToDead(); 
		bool isTidInActiveQueue(int tid); 
		bool isTidInActiveQueue(); 
		bool isTidInBlockingQueue(); 
		bool isTidInQueue(); 
		bool findPthreadIDInActiveAndUpdate(int tid); 
		void insertPthreadIDAfterActiveTail(); 
		void insertTidAfterActiveTail(int tid); 
		void insertAfterActiveTail(int tid); 
		void insertAfterActiveTail(); 
		void insertBeforeActiveHead(); 
		void insertInvalidBeforeActiveHead(); 
		void insertAfterBlockingTail(uint64_t syncVarAddr, uint64_t secondSyncVarAddr = 0); 
		int getActiveHead(); 
		std::string PrintTids(); 
		int activateTids(uint64_t syncVarAddr, int activateNum); 
		bool isActiveHead(); 
		bool isActiveHeadNoAssert(); 
		bool isBlockingHead(); 
		bool isSyncVarInBlockingQueue(uint64_t syncVarAddr); 
		bool isBlockingHeadWithSyncVar(uint64_t syncVarAddr); 
		void delFromActiveHead(); 
		int delFromBlockingHeadWithSyncVar(uint64_t syncVarAddr);
		void delFromBlockingQueue(BlockingTidInfo *bti);
		void insertToBlockingHead(BlockingTidInfo *bti);
		BlockingTidInfo *getActivatedTidInfo(uint64_t syncVarAddr);
		void insertToActiveHead(int tid);
		void delFromQueue(); 
		int getInTaskThreadNum();
		void setInTaskThreadNum(int newNum);
		int getOutTaskThreadNum();
		void setOutTaskThreadNum(int newNum);
		int getTidMap();
		int getTidMap(int tid);
                int getNextTid();
		bool findTidMap(int tid);
		void incTidMap();
		void updateTidMapForFork(int parentTid, int childTid);
		void delTidFromMap(int tid);
	};

	// The module to enforce round robin order of landmarks and record the order.
	class Memoizer: public TernSched {
	private:
		Int_To_Int_Map isInTask;
		Long_To_Int_Map mutexLocked;
		Long_To_Int_Map mutexOwner;
		Int_To_Cond_Var_Map condVarForMutex;
		Long_To_Int_Map rwLocked;
		Int_To_Cond_Var_Map condVarForRWLock;
		vector<char *> windowTidTotalOrder;
		pthread_cond_t condVarForActiveTurn;
		pthread_cond_t condVarForBlockingTurn;

		// Controlling window order.
		std::string windowConstraintString;
		int globalWindowSerialNum;
		Int_To_Int_Map mmzEnterTaskTidOrder;

		// Queue to enforce order.
		class Queue *queue;
		struct LogMemAccess memAccess;
		SchedCache *schedCache;

		// Branch slicing module.
		struct OmitBranch *oBr;
		Long_To_Int_Map brInstOmitted;

		// Debug utility.
		vector <char *> taskConcreteBytes;

	protected:
 		void appendTidTotalOrder(const char *funcName, int myTid, int instId, const uint64_t *args, KInstruction *ki, int hookId);
		void printMemAccess();
		void appendLandmarkOp(const char *syncEventName, int log_type, const uint64_t *args);
		void dumpTidsToFile();
		void appendWindowString(std::string str);
		bool mutexAvailable(uint64_t mutexAddr);
		void setMutexLocked(uint64_t mutexAddr);
		void setMutexUnLocked(uint64_t mutexAddr);
		void setReadLocked(uint64_t rwLockAddr);
		void setWriteLocked(uint64_t rwLockAddr);
		void setRWUnLocked(uint64_t rwLockAddr);
		bool isReadLocked(uint64_t rwLockAddr);
		bool isWriteLocked(uint64_t rwLockAddr);
		bool isRWUnLocked(uint64_t rwLockAddr);
		void initCondVarForMutex(uint64_t mutexAddr);
		void initCondVarForRWLock(uint64_t rwLockAddr);
		int waitForBlockingTurn(uint64_t syncVarAddr);

		// Before landmark calls.
		int beforePthreadMutexLock(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int beforePthreadMutexUnLock(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int beforePthreadCreate(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int beforeAcceptSocket(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int beforeSleep(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int beforeSignalWait(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int beforePthreadJoin(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int beforePthreadCondSignal(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int beforePthreadCondBroadcast(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int beforePthreadCondWait(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int beforePthreadCondTimedWait(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int beforePthreadBarrierWait(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		void ternBeforeMakeSymbolic(const char *cacheName, KInstruction *ki, int instId, int hookId);
		int beforePthreadRWLock(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, bool isWriteLock, int hookId);
		int beforePthreadRWUnLock(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);

		// After landmark calls.
		int afterPthreadMutexLock(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int afterPthreadMutexUnLock(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int afterPthreadCreate(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int afterAcceptSocket(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int afterSleep(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int afterSignalWait(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int afterPthreadJoin(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int afterPthreadCondSignal(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int afterPthreadCondBroadcast(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int afterPthreadCondWait(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int afterPthreadCondTimedWait(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		int afterPthreadBarrierWait(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);
		void ternAfterMakeSymbolic();
		int afterPthreadRWLock(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, bool isWriteLock, int hookId);
		int afterPthreadRWUnLock(const char *funcName, const uint64_t *args, KInstruction *ki, int instId, int hookId);

	public:
		/* Below are the interface provided to the KLEE engine. */
		Memoizer(SchedCache *sc);

		/* Interface inheritated from TernSched class. */
		virtual void appStart();
		virtual void appEnd(int *info);

		virtual void taskStart(int *winNum, int instId, int hookId, void *addr, unsigned nBytes, const char *symVarName);
		virtual int taskEnd(int *winNum, int instId, int hookId);

		virtual void threadCreateStart(unsigned arg);
		virtual void threadCreateEnd(unsigned arg, int pthreadId);
		virtual void threadExit(int *info, int hookId, bool realExit = 0);

		/* Heming: I have already added an interpreting mechanism to the KLEE interpreter (SpecialFunctionHandler)
			with the name "tern_symbolic", so here I just define this function body as empty. */
		virtual void symbolic(void *addr, unsigned nbytes, const char *name) {}

		virtual void beforeLandmark(unsigned arg0, unsigned hook_id, int instId, uint64_t arg1, uint64_t arg2, uint64_t arg3/*, std::string callStackStr*/);
		virtual void afterLandmark(pthread_t *tid, unsigned arg0, unsigned hook_id, int instId, uint64_t ret_val, uint64_t arg1, uint64_t arg2, uint64_t arg3/*, std::string callStackStr*/);

		virtual void waitForTurn(const char *landmarkName, unsigned arg, int hookId = -1);
		virtual void giveUpTurn(const char *landmarkName, unsigned arg);

		/* Utility functions. */
		int ternMyTid(int tid);
		void recordMemAccess(bool isLocal, ExecutionState &state, KInstruction *ki, int addr, int instId);
		Queue *getQueue() {return queue;};
		std::string getCallStackStr(ExecutionState &state);
		bool threadInTask();
		void updateTidMapForFork(int parentTid, int childTid);
		void initBranchSlicingModule(Module *m);
		bool omitBranch(BranchInst *branch, ref<Expr> condition);
		int mmzEnterTaskOrder();
		int mmzTaskConcreteBytes(void *addr);
	};
}

#endif

